import torch
from TransformerEncoder import TransformerEncoder
from TransformerDecoder import TransformerDecoder
from EncoderDecoder import EncoderDecoder

if __name__ == '__main__':
    enc_inputs = torch.randint(0, 10, (20, 12))  # (batch_size, seq_len)
    enc_valid_lens = torch.randint(1, 6, (20,))  # valid_lens (<pad>)
    dec_inputs = torch.randint(0, 22, (20, 11))  # (batch_size, seq_len)

    encoder = TransformerEncoder(10, 512, 2, 512,
                                 512, 8, 512, 0.5)
    decoder = TransformerDecoder(22, 512, 2, 512,
                                 512, 8, 512, 0.5)

    transformer = EncoderDecoder(encoder, decoder)
    outputs, state = transformer(enc_inputs, dec_inputs, enc_valid_lens)
    print(outputs.shape)
